# coding: UTF-8
import argparse
from importlib import import_module

import numpy as np
import torch

from train_eval import test
from utils import build_iterator, load_dataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

dataset = 'TikTok'  # 数据集
parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, required=True,
                    help='choose a model: Bert, ERNIE')
args = parser.parse_args()

model_name = args.model  # bert
x = import_module(f'models.{model_name}')
config = x.Config(dataset)
np.random.seed(3407)
torch.manual_seed(3407)
torch.cuda.manual_seed_all(3407)
torch.backends.cudnn.deterministic = True  # 保证每次结果一样

# infer
test_data = load_dataset(config, config.test_path, config.pad_size)
test_iter = build_iterator(test_data, config)

model = x.Model(config).to(device)
model.eval()

predicted_labels = test(config, model, test_iter, verbose=0)

# label_mapping = {}
# with open("/home/hwxu/Projects/Research/asset/TikTok/data/no_stopwords/class.txt", "r", encoding="utf-8") as file:
#     for label_id, label_name in enumerate(file):
#         label_mapping[int(label_id)] = label_name.replace('\n', '')

# predicted_labels = [label_mapping[label] for label in predicted_labels]
# print(predicted_labels)
